package simpledb.execution;

import simpledb.common.Type;
import simpledb.storage.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Knows how to compute some aggregate over a set of StringFields.
 */
public class StringAggregator implements Aggregator {

    private static final long serialVersionUID = 1L;

    private final int gbfield;
    private final Type gbfieldtype;
    private final int afield;
    private final Op what;
    // 存储分组到聚合状态的映射：key为分组字段值（无分组时为null），value为聚合中间结果
    private Map<Field, Integer> groupMap;

    /**
     * Aggregate constructor
     * @param gbfield the 0-based index of the group-by field in the tuple, or NO_GROUPING if there is no grouping
     * @param gbfieldtype the type of the group by field (e.g., Type.INT_TYPE), or null if there is no grouping
     * @param afield the 0-based index of the aggregate field in the tuple
     * @param what aggregation operator to use -- only supports COUNT
     * @throws IllegalArgumentException if what != COUNT
     */

    public StringAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        // some code goes here
        // 检查是否是支持的操作
        if (what != Op.COUNT) {
            throw new IllegalArgumentException("StringAggregator only supports COUNT operation");
        }
        this.gbfield = gbfield;
        this.gbfieldtype = gbfieldtype;
        this.afield = afield;
        this.what = what;
        // 使用hashmap处理无分组key为null的情况
        this.groupMap = new HashMap<>();
    }

    /**
     * Merge a new tuple into the aggregate, grouping as indicated in the constructor
     * @param tup the Tuple containing an aggregate field and a group-by field
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        // some code goes here
        Field groupKey;
        if (gbfield == Aggregator.NO_GROUPING) {
            groupKey = null;
        }else {
            groupKey = tup.getField(gbfield);
        }

        groupMap.put(groupKey, groupMap.getOrDefault(groupKey, 0)+1);
        // System.out.println(groupKey + ": " + groupMap.get(groupKey));
    }

    /**
     * Create a OpIterator over group aggregate results.
     *
     * @return a OpIterator whose tuples are the pair (groupVal,
     *   aggregateVal) if using group, or a single (aggregateVal) if no
     *   grouping. The aggregateVal is determined by the type of
     *   aggregate specified in the constructor.
     *   有TupleIterator可以用，就不自己实现类了，但是它没有对iterator为null的时候的判断
     */
    public OpIterator iterator() {
        // some code goes here
        // throw new UnsupportedOperationException("please implement me for lab2");
        TupleDesc td;
        List<Tuple> resultTuples = new ArrayList<>();
        if (gbfield==Aggregator.NO_GROUPING) {
            td = new TupleDesc(new Type[]{Type.INT_TYPE});
            Tuple tuple = new Tuple(td);
            tuple.setField(0, new IntField(groupMap.get(null)));
            resultTuples.add(tuple);
        }else {
            td = new TupleDesc(new Type[]{gbfieldtype, Type.INT_TYPE});
            int index = 0;
            for (Map.Entry<Field, Integer> entry : groupMap.entrySet()) {
                Tuple tuple = new Tuple(td);
                tuple.setField(0, entry.getKey());
                tuple.setField(1, new IntField(groupMap.get(entry.getKey())));
                resultTuples.add(tuple);
            }
        }
        return new TupleIterator(td, resultTuples);
    }

}
